import numpy as np
import pandas as pd
from sklearn import metrics
import lightgbm as lgb
from sklearn.model_selection import StratifiedKFold
import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'


trn_data = pd.read_csv('./data/trn_per_100.csv', header=0)
X = trn_data.iloc[:, :-2]
y = trn_data['type']
features = [x for x in X.columns if x not in ['ship', 'type', 'time', 'diff_time', 'date']]
fold = StratifiedKFold(n_splits=5, shuffle=True, random_state=100)

params = {
    'n_estimators': 5000,
    'boosting_type': 'gbdt',
    'objective': 'multiclass',
    'num_class': 3,
    'early_stopping_rounds': 100
}

models = []
oof = np.zeros((len(X), 3))
for index, (train_idx, val_idx) in enumerate(fold.split(X, y)):
    train_set = lgb.Dataset(X.iloc[train_idx], y.iloc[train_idx])
    val_set = lgb.Dataset(X.iloc[val_idx], y.iloc[val_idx])

    model = lgb.train(params, train_set, valid_sets=[train_set, val_set], verbose_eval=100)
    models.append(model)
    val_pred = model.predict(X.iloc[val_idx])
    oof[val_idx] = val_pred
    val_y = y.iloc[val_idx]
    val_pred = np.argmax(val_pred, axis=1)
    print(index, 'val f1', metrics.f1_score(val_y, val_pred, average='macro'))


oof = np.argmax(oof, axis=1)
print('oof f1', metrics.f1_score(oof, y, average='macro'))
# 0.8701544575329372

import matplotlib.pyplot as plt
header = pd.DataFrame(X.columns)
feat_imp = model.feature_importance()
sort_ind = feat_imp.argsort()
A = pd.concat([header.ix[sort_ind, :].reset_index().drop('index', axis=1), pd.DataFrame(feat_imp[sort_ind])], axis=1, ignore_index=True)
plt.figure(figsize=(10, 10))
# plt.barh(A[0], A[1], height=0.5)
# plt.savefig('fts/gbdt_fea_imp.png')
# plt.show()

